#include "lv_jpeg_decoder.h"
#include "lvgl.h"
#include "string.h"
#include "HWJpegDecoder.h"
#include "jpeg.h"
#include "DMA2D.h"

static lv_res_t decoder_info(struct _lv_img_decoder_t *decoder, const void *src, lv_img_header_t *header);
static lv_res_t decoder_open(lv_img_decoder_t *decoder, lv_img_decoder_dsc_t *dsc);
static void decoder_close(lv_img_decoder_t *decoder, lv_img_decoder_dsc_t *dsc);

static void YCbCrToRGB(HWJpegDecoder::ImageInfo &info, const void *data, uint32_t line, uint32_t numOfLine);

static HWJpegDecoder *jpegDecoder = nullptr;
static uint8_t *globalImageData = nullptr;

void lv_jpeg_decoder_init(void)
{
    lv_img_decoder_t *dec = lv_img_decoder_create();
    lv_img_decoder_set_info_cb(dec, decoder_info);
    lv_img_decoder_set_open_cb(dec, decoder_open);
    lv_img_decoder_set_close_cb(dec, decoder_close);

    dec->user_data = nullptr;
}

static lv_res_t decoder_info(struct _lv_img_decoder_t *decoder, const void *src, lv_img_header_t *header)
{
    lv_img_src_t src_type = lv_img_src_get_type(src); /*Get the source type*/

    /*If it's a JPEG file...*/
    if (src_type == LV_IMG_SRC_FILE)
    {
        const char *fn = (const char *)src;
        const char *ext = lv_fs_get_ext(fn);
        if (strcmp(ext, "jpg") == 0 || strcmp(ext, "jpeg") == 0 || strcmp(ext, "JPG") == 0 || strcmp(ext, "JPEG") == 0)
        { /*Check the extension*/

            if (decoder->user_data == nullptr)
            {
                HWJpegDecoder *jpegDecoder = new HWJpegDecoder(&hjpeg);
                jpegDecoder->setOutputCallback(YCbCrToRGB);
                decoder->user_data = jpegDecoder;
            }

            HWJpegDecoder *jpegDecoder = (HWJpegDecoder *)decoder->user_data;
            jpegDecoder->reset();
            jpegDecoder->setFullInput(false);
            jpegDecoder->setInputDataSource(fn);
            if (jpegDecoder->decode(true) != HWJpegDecoder::NoData)
            {
                jpegDecoder->reset();
                return LV_RES_INV;
            }

            HWJpegDecoder::ImageInfo &info = jpegDecoder->getImageInfo();

            /*Save the data in the header*/
            header->always_zero = 0;
            header->cf = LV_IMG_CF_TRUE_COLOR;
            header->w = info.width;
            header->h = info.height;

            return LV_RES_OK;
        }
    }
    /*If it's a PNG file in a  C array...*/
    else if (src_type == LV_IMG_SRC_VARIABLE)
    {
        if (decoder->user_data == nullptr)
        {
            HWJpegDecoder *jpegDecoder = new HWJpegDecoder(&hjpeg);
            jpegDecoder->setOutputCallback(YCbCrToRGB);
            decoder->user_data = jpegDecoder;
        }

        const lv_img_dsc_t *img_dsc = (const lv_img_dsc_t *)src;
        HWJpegDecoder *jpegDecoder = (HWJpegDecoder *)decoder->user_data;
        jpegDecoder->reset();
        jpegDecoder->setFullInput(false);
        jpegDecoder->setInputDataSource(img_dsc->data, img_dsc->data_size);
        if (jpegDecoder->decode(true) != HWJpegDecoder::NoData)
        {
            jpegDecoder->reset();
            return LV_RES_INV;
        }

        HWJpegDecoder::ImageInfo &info = jpegDecoder->getImageInfo();

        header->always_zero = 0;
        header->cf = LV_IMG_CF_TRUE_COLOR;
        header->w = info.width;
        header->h = info.height;
        return LV_RES_OK;
    }

    return LV_RES_INV; /*If didn't succeeded earlier then it's an error*/
}

static lv_res_t decoder_open(lv_img_decoder_t *decoder, lv_img_decoder_dsc_t *dsc)
{
    uint8_t *img_data = nullptr;

    /*If it's a JPEG file...*/
    if (dsc->src_type == LV_IMG_SRC_FILE)
    {
        const char *fn = (const char *)dsc->src;
        const char *ext = lv_fs_get_ext(fn);
        if (strcmp(ext, "jpg") == 0 || strcmp(ext, "jpeg") == 0 || strcmp(ext, "JPG") == 0 || strcmp(ext, "JPEG") == 0)
        { /*Check the extension*/

            if (decoder->user_data == nullptr)
            {
                HWJpegDecoder *jpegDecoder = new HWJpegDecoder(&hjpeg);
                jpegDecoder->setOutputCallback(YCbCrToRGB);
                decoder->user_data = jpegDecoder;
            }

            HWJpegDecoder *jpegDecoder = (HWJpegDecoder *)decoder->user_data;
            jpegDecoder->reset();
            jpegDecoder->setFullInput(false);
            jpegDecoder->setInputDataSource(fn);
            if (jpegDecoder->decode(true) != HWJpegDecoder::NoData)
            {
                jpegDecoder->reset();
                return LV_RES_INV;
            }

            HWJpegDecoder::ImageInfo &info = jpegDecoder->getImageInfo();
            img_data = new uint8_t[info.width * info.height * sizeof(lv_color_t)];
            if (img_data == nullptr)
            {
                LV_LOG_WARN("out of memory\r\n");
                return LV_RES_INV;
            }

#if DECODE_FULLINPUT
            jpegDecoder->setFullInput(true);
#endif

            globalImageData = img_data;
            if (jpegDecoder->decode(false) != HWJpegDecoder::NoError)
            {
                jpegDecoder->reset();
                return LV_RES_INV;
            }

            jpegDecoder->waitDecodeFinish();

            /*Convert the image to the system's color depth*/
            // convert_color_depth(img_data, info.width * info.height);
            dsc->img_data = img_data;
            return LV_RES_OK; /*The image is fully decoded. Return with its pointer*/
        }
    }
    /*If it's a PNG file in a  C array...*/
    else if (dsc->src_type == LV_IMG_SRC_VARIABLE)
    {
        const lv_img_dsc_t *img_dsc = (const lv_img_dsc_t *)dsc->src;

        if (decoder->user_data == nullptr)
        {
            HWJpegDecoder *jpegDecoder = new HWJpegDecoder(&hjpeg);
            jpegDecoder->setOutputCallback(YCbCrToRGB);
            decoder->user_data = jpegDecoder;
        }

        HWJpegDecoder *jpegDecoder = (HWJpegDecoder *)decoder->user_data;
        jpegDecoder->reset();
        jpegDecoder->setFullInput(false);
        jpegDecoder->setInputDataSource(img_dsc->data, img_dsc->data_size);
        if (jpegDecoder->decode(true) != HWJpegDecoder::NoData)
        {
            jpegDecoder->reset();
            return LV_RES_INV;
        }

        HWJpegDecoder::ImageInfo &info = jpegDecoder->getImageInfo();
        img_data = new uint8_t[info.width * info.height * sizeof(lv_color_t)];
        if (img_data == nullptr)
        {
            LV_LOG_WARN("out of memory\r\n");
            return LV_RES_INV;
        }

#if DECODE_FULLINPUT
            jpegDecoder->setFullInput(true);
#endif

        globalImageData = img_data;
        if (jpegDecoder->decode(false) != HWJpegDecoder::NoError)
        {
            jpegDecoder->reset();
            return LV_RES_INV;
        }

        jpegDecoder->waitDecodeFinish();

        /*Convert the image to the system's color depth*/
        // convert_color_depth(img_data, info.width * info.height);

        dsc->img_data = img_data;
        return LV_RES_OK; /*Return with its pointer*/
    }

    return LV_RES_INV; /*If not returned earlier then it failed*/
}

static void decoder_close(lv_img_decoder_t *decoder, lv_img_decoder_dsc_t *dsc)
{
    (void)decoder;

    if (dsc->img_data)
    {
        delete[] (uint8_t *)dsc->img_data;
        dsc->img_data = NULL;
    }
}

static void YCbCrToRGB(HWJpegDecoder::ImageInfo &info, const void *data, uint32_t line, uint32_t numOfLine)
{
    DMA2DParam param;

    DMA2DParamSetSrc(param, (uint32_t)data);
    DMA2DParamSetDst(param, (uint32_t)globalImageData);
    DMA2DParamSetX(param, 0);
    DMA2DParamSetY(param, line * info.MCUHeight);
    DMA2DParamSetXSize(param, info.width);
    DMA2DParamSetYSize(param, numOfLine * info.MCUHeight);
    DMA2DParamSetChromaSampling(param, info.chromaSubsampling);

#if LV_COLOR_DEPTH == 16
    DMA2DParamSetOutputPixelFormat(param, LTDC_PIXEL_FORMAT_RGB565);
#elif LV_COLOR_DEPTH == 32
    DMA2DParamSetOutputPixelFormat(param, LTDC_PIXEL_FORMAT_ARGB8888);
#else
#warning "no support for this color depth"
#endif

    DMA2D_CopyYCbCrToRGB(param);
}